#_"SPDX-License-Identifier: GPL-3.0"

(ns ^:no-doc sicmutils.polynomial.exponent
  "This namespace provides an implementation of a sparse representation of the
  exponent portion of a term of a polynomial, sometimes called a 'monomial'."
  (:refer-clojure :exclude [empty assoc])
  (:require [clojure.core :as core]
            [sicmutils.util :as u]))

;; ## Sparse Monomial Exponents
;;
;; This namespace is used by `sicmutils.polynomial.impl` and
;; `sicmutils.polynomial` to implement a full polynomial data structure.
;;
;; Polynomials are sums of monomial terms; a monomial is a pair of some non-zero
;; coefficient and a product of some number of variables, each raised to some
;; power.
;;
;; We represent the exponents of a monomial with an ordered mapping of variable
;; index => the exponent of that variable. $x^2z^3$ is represented as `{0 2, 2
;; 3}`, for example. Polynomials are linear combinations of the exponents.
;;
;; ## Constructors

(def ^{:doc "Accepts alternating pairs of integers representing <indeterminate
 index>, <exponent value> and returns a `sorted-map` representing the exponent
 portion of a polynomial term."
       :arglists '([& i-expt-pairs])}
  make
  #'sorted-map)

(def ^{:doc "Singleton instance of an empty exponent map."}
  empty
  (make))

(defn dense->exponents
  "Accepts a sequence of pairs of indeterminate index => power, and returns a
  sparse representation the exponents portion of a monomial.

  For example:

  ```clojure
  (dense->exponents [1 4 0 0 2])
  ;;=> {0 1, 1 4, 4 2}
  ```"
  [idx->pow]
  (reduce-kv (fn [acc i x]
               (if (zero? x)
                 acc
                 (core/assoc acc i x)))
             empty
             idx->pow))

(defn mul
  "Returns a new exponent vector generated by multiplying all terms in exponent
  vectors `l` and `r`."
  [l r]
  (merge-with + l r))

(defn div
  "Returns a new exponent vector generated by dividing terms in exponent vector
  `l` by terms in `r`."
  [l r]
  (dense->exponents
   (merge-with + l (u/map-vals - r))))

(defn gcd
  "Returns the exponent vector that is the greatest common divisor of the exponent
  vectors `l` and `r`.

  Calling [[gcd]] with a single argument acts as identity."
  ([l] l)
  ([l r]
   (let [l' (select-keys l (keys r))
         r' (select-keys r (keys l))]
     (merge-with min l' r'))))

(defn lcm
  "Returns the exponent vector that is the least common multiple of the exponent
  vectors `l` and `r`.

  Calling [[lcm]] with no arguments returns the empty exponent vector; calling
  with a single argument acts as identity."
  ([] empty)
  ([l] l)
  ([l r]
   (merge-with max l r)))

(defn every-power?
  "Returns true if `f` returns true for every positive non-zero exponent in the
  supplied exponent vector `m`, false otherwise.

  Defaults to `true` if `m` is empty."
  [f m]
  (every? f (vals m)))

(defn assoc
  "Replaces the entry for variable `x` in the exponents vector `m` with power `n`.
  If `n` is zero, removes `x`'s entry from the returned exponent vector."
  [m x n]
  (if (zero? n)
    (dissoc m x)
    (core/assoc m x n)))

(defn lower
  "Given some exponent vector `expts`, and an optional variable index
  `i` (defaults to `0`), returns a new exponent vector with:

  - the `i`th variable's entry removed
  - decremented index for each variable with index > `i`.

  For example:

  ```clojure
  (lower {0 3, 1 2, 4 4})
  ;;=> {0 2, 3 4}

  (lower {0 3, 1 2, 4 4} 1)
  ;;=> {0 3, 3 4}
  ```"
  ([expts]
   (lower expts 0))
  ([expts i]
   (reduce-kv (fn [acc k v]
                (if (> k i)
                  (core/assoc acc (dec k) v)
                  (core/assoc acc k v)))
              empty
              (dissoc expts i))))

(defn raise
  "Given some exponent vector `expts`, an optional variable index `i` (defaults to
  `0`) and an optional exponent power `n` (defaults to `0`), returns a new
  exponent vector with:

  - incremented indices for each variable with index >= `i`
  - a new `i`th variable created with power `n`

  For example:

  ```clojure
  (raise {0 3, 1 2, 4 4})
  ;;=> {1 3, 2 2, 5 4}

  (raise {0 3, 1 2, 4 4} 1)
  ;;=> {0 3, 2 2, 5 4}

  (raise {0 3, 1 2, 4 4} 1)
  ;;=> {0 3, 1 10, 2 2, 5 4}
  ```"
  ([expts]
   (raise expts 0 0))
  ([expts i]
   (raise expts i 0))
  ([expts i n]
   (let [m (reduce-kv (fn [acc k v]
                        (if (>= k i)
                          (core/assoc acc (inc k) v)
                          (core/assoc acc k v)))
                      empty
                      expts)]
     (if (zero? n)
       m
       (core/assoc m i n)))))

(defn monomial-degree
  "Returns the [monomial degree](https://en.wikipedia.org/wiki/Monomial#Degree) of
  the exponent vector `m`, ie, the sum of the powers of all variables in `m`.

  If the optional `i` is supplied, returns the degree of the `i`th variable, ie,
  the entry for `i` in `m`, defaulting to `0`."
  ([m]
   (apply + (vals m)))
  ([m i]
   (m i 0)))

(defn ->sort+unsort
  "Given a power product `m`, returns a pair of `sort` and `unsort` functions of a
  single power product argument.

  `sort` rearranges the indices of its argument to match the order of increasing
  variable degree in `m`. `unsort` undoes this transformation."
  [m]
  (let [indices (range (count m))
        order (into [] (sort-by m (keys m)))]
    (letfn [(sort [m']
              (into empty
                    (mapcat (fn [i]
                              (when-let [v (m' (order i))]
                                [[i v]])))
                    indices))
            (unsort [m']
              (into empty
                    (mapcat (fn [i]
                              (when-let [v (m' i)]
                                [[(order i) v]])))
                    indices))]
      [sort unsort])))

;; ## Monomial Orderings
;;
;; This section implements a number of [Monomial
;; orderings](https://en.wikipedia.org/wiki/Monomial_order) that are useful for
;; making various polynomial algorithms efficient.
;;
;; Each of the following functions matches Java's "comparator" interface:

;; - If the left argument is greater than the right argument, the function
;;   returns a negative number
;; - equal values return `0`
;; - if the right value is greater, returns a positive number

(defn lex-order
  "Comparator that responds based on the [lexicographic
  order](https://en.wikipedia.org/wiki/Monomial_order#Lexicographic_order), or
  'lex order', of the exponent vectors `xs` and `ys`. Accepts any sequence of
  pairs of the form `[variable, power]` for `xs` and `ys`.

  Lex order first compares the power of variable `0`, then, in case of equality,
  variable `1` and so on. If all powers match, returns `0`.

  If `:reverse? true` is passed, the inputs are compared in reverse lexicographic order.

  Reverse here means two things:

  1. the variables are considered in reverse order; the variable with the
     largest index is compared first, then the next-largest, and on down the line.
  2. The _smaller_ exponent is grevlex-greater in this comparison."
  ([xs ys & {:keys [reverse?]}]
   (let [xs (vec (if reverse? (rseq xs) xs))
         ys (vec (if reverse? (rseq ys) ys))]
     (loop [i (long 0)]
       (let [x (nth xs i nil)
             y (nth ys i nil)]
         (cond (and (not x) (not y)) 0
               (not x) -1
               (not y)  1
               :else (let [bit (compare (nth x 0) (nth y 0))]
                       (cond (zero? bit)
                             (let [xv (nth x 1)
                                   yv (nth y 1)]
                               (if (= xv yv)
                                 (recur (inc i))
                                 (if reverse?
                                   (- yv xv)
                                   (- xv yv))))
                             (neg? bit) 1
                             :else -1))))))))

(defn graded-lex-order
  "Comparator that responds based on the [graded lexicographic
  order](https://en.wikipedia.org/wiki/Monomial_order#Graded_lexicographic_order),
  or 'grlex order', of the exponent vectors `xs` and `ys`. Accepts any sequence
  of pairs of the form `[variable, power]` for `xs` and `ys`.

  grlex order first compares the total degree of `xs` and `ys`, and falls back
  to [[lex-order]] in case of a tie. See [[lex-order]] for details on this
  case."
  [xs ys]
  (let [xd (monomial-degree xs)
        yd (monomial-degree ys)]
    (if (= xd yd)
      (lex-order xs ys)
      (- xd yd))))

(defn graded-reverse-lex-order
  "Comparator that responds based on the [graded reverse lexicographic
  order](https://en.wikipedia.org/wiki/Monomial_order#Graded_reverse_lexicographic_order),
  or 'grevlex order', of the exponent vectors `xs` and `ys`. Accepts any
  sequence of pairs of the form `[variable, power]` for `xs` and `ys`.

  grevlex order first compares the total degree of `xs` and `ys`, and falls back
  to reverse [[lex-order]] in case of a tie. Reverse here means two things:

  1. the variables are considered in reverse order; the variable with the
     largest index is compared first, then the next-largest, and on down the line.
  2. The _smaller_ exponent is grevlex-greater in this comparison."
  [xs ys]
  (let [xd (monomial-degree xs)
        yd (monomial-degree ys)]
    (if (= xd yd)
      (lex-order xs ys :reverse? true)
      (- xd yd))))
